{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparison of the results of the stable diffusion optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from optimum.intel import OVStableDiffusionPipeline\n",
    "from diffusers.training_utils import set_seed\n",
    "from IPython.display import display"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the original pipeline\n",
    "This pipeline was fine-tuned on the public [dataset](https://huggingface.co/datasets/lambdalabs/pokemon-blip-captions) with Pokemon images and the correspoinding captions. You can find the source model and the description [here](https://huggingface.co/svjack/Stable-Diffusion-Pokemon-en)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = OVStableDiffusionPipeline.from_pretrained(\"OpenVINO/stable-diffusion-pokemons-fp32\", compile=False)\n",
    "pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)\n",
    "\n",
    "pipe.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's fix the seed for reproducibility.\n",
    "set_seed(42)\n",
    "\n",
    "prompt = \"cartoon bird\"\n",
    "output = pipe(prompt, num_inference_steps=50, output_type=\"pil\")\n",
    "display(output.images[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the quantized pipeline\n",
    "Now we run the quantized pipeline that was obtained with Quantization-Aware Training on the same dataset. The original model was used as a baseline for quantization. The resulted model can be found [here](https://huggingface.co/OpenVINO/Stable-Diffusion-Pokemon-en-quantized)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quantized_pipe = OVStableDiffusionPipeline.from_pretrained(\"OpenVINO/Stable-Diffusion-Pokemon-en-quantized\", compile=False)\n",
    "quantized_pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)\n",
    "quantized_pipe.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the same seed to compare\n",
    "set_seed(42)\n",
    "\n",
    "output = quantized_pipe(prompt, num_inference_steps=50, output_type=\"pil\")\n",
    "display(output.images[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the highly optimized pipeline\n",
    "Here, we run the pipeline optimized with a combination of Token Merging Method and Quantization-aware training. The resulted model can be found [here](https://huggingface.co/OpenVINO/stable-diffusion-pokemons-tome-quantized)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_pipe = OVStableDiffusionPipeline.from_pretrained(\"OpenVINO/stable-diffusion-pokemons-tome-quantized\", compile=False)\n",
    "optimized_pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)\n",
    "optimized_pipe.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the same seed to compare\n",
    "set_seed(42)\n",
    "\n",
    "output = optimized_pipe(prompt, num_inference_steps=50, output_type=\"pil\")\n",
    "display(output.images[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can see the difference of the difference in the results and the time required to generate the image."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('stable_diffusion')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "7918409a64d3d4275e0103fc4443d9be5863d1df136c02ed032407c7ae821339"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
